import os
import argparse
import pandas as pd
import freesasa
from Bio.PDB import PDBParser
from tqdm import tqdm


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset_input", type=str, default=None)
    parser.add_argument("--pdb_file", type=str, default=None)
    parser.add_argument("--out_file", type=str, default="sasa.csv")
    parser.add_argument("--type", type=str, choices=["residue", "protein"], default="protein")
    args = parser.parse_args()
    
    if args.dataset_input is not None:
        out_info = {"pdb": [], "sasa": []}

        pdbs = sorted(os.listdir(args.dataset_input))
        for pdb in tqdm(pdbs):
            pdb_file = os.path.join(args.dataset_input, pdb)
            
            pdb_parser = PDBParser()
            structure = pdb_parser.get_structure('PDB_ID', pdb_file)
            model = structure[0]

            seq_len = 0
            for chain in model:
                for residue in chain:
                    if residue.id[0] == ' ':
                        seq_len += 1
            sasa_result, sasa_classes = freesasa.calcBioPDB(structure)
            avg_sasa = round(sasa_result.totalArea() / seq_len, 3)
            
            out_info["pdb"].append(pdb)
            out_info["sasa"].append(avg_sasa)
        pd.DataFrame(out_info).to_csv(args.out_file, index=False)
    else:
        pdb_parser = PDBParser()
        structure = pdb_parser.get_structure('PDB_ID', args.pdb_file)
        sasa_result, sasa_classes = freesasa.calcBioPDB(structure)
        avg_sasa = round(sasa_result.totalArea() / seq_len, 3)

        print(avg_sasa)